# -*- coding: utf-8 -*-

import os.path
import cv2
import gzip
import struct
import numpy

image_dir = '/Users/vista/Documents/PIC/ocr/data'

data_dir = '/Users/vista/PycharmProjects/data/ocr_data'

images_file_name_digit = 'digit-images-ubyte'

labels_file_name_digit = 'digit-labels-ubyte'

images_file_name_letter = 'letter-images-ubyte'

labels_file_name_letter = 'letter-labels-ubyte'

def _read32(bytestream):
  dt = numpy.dtype(numpy.uint32).newbyteorder('>')
  #return numpy.frombuffer(bytestream.read(4), dtype=dt) old ver.
  return numpy.frombuffer(bytestream.read(4), dtype=dt)[0]

def list_subdir(path):

    list = []

    if os.path.exists(path):

        files = os.listdir(path)
        for file in files:

            m = os.path.join(path,file)

            if os.path.isdir(m):
                h = os.path.split(m)

                list.append(h[1])
    return list

def append_into(arr,char,AA):
    dir = os.path.join(image_dir, char)

    files = os.listdir(dir)

    for file in files:

        obj = AA()

        obj.char = char

        full_path = os.path.join(dir, file)

        if full_path.find('.jpg') > 0:
            img = cv2.imread(full_path, cv2.IMREAD_GRAYSCALE)

            obj.img = img

            arr.append(obj)

def read_all_images(chars):

    class AA:
        pass

    digit = []
    letter = []

    for char in chars:

        if '0' <= char <= '9':
            append_into(digit, char , AA)
        else:
            append_into(letter, char, AA)

    return digit , letter

def write_file(all,type):

    # image_bin_file_path = ''
    # label_bin_file_path = ''

    if type == 0:
        image_bin_file_path = os.path.join(data_dir,images_file_name_digit)
        label_bin_file_path = os.path.join(data_dir,labels_file_name_digit)
    else:
        image_bin_file_path = os.path.join(data_dir,images_file_name_letter)
        label_bin_file_path = os.path.join(data_dir,labels_file_name_letter)

    image_bin_file = open(image_bin_file_path,'wb')
    label_bin_file = open(label_bin_file_path,"wb")

    # write info about images

    bytes = struct.pack('iiii',2052,len(all),56,28)

    image_bin_file.write(bytes)

    for sub in all:
        img = sub.img

        # image_bin_file.write( struct.pack(str(img.shape[0]*img.shape[1])+'B',img[:]) )

        for r in range(img.shape[0]):
            for c in range(img.shape[1]):
                # image_bin_file.write(struct.pack('B',img[r,c]))
                image_bin_file.write(img[r,c])

    image_bin_file.close()


    bytes = struct.pack('ii',2050,len(all))

    label_bin_file.write(bytes)

    for sub in all:
        c = sub.char[0]

        if ord('0') <= ord(c) <= ord('9'):
            c = ord(c) - ord('0')
        else:
            c = ord(c) - ord('A')

        label_bin_file.write( struct.pack('B',c) )

    label_bin_file.close()

def open_file(type):

    image_bin_file_path = ''
    label_bin_file_path = ''

    if type == 0:
        image_bin_file_path = os.path.join(data_dir,images_file_name_digit)
        label_bin_file_path = os.path.join(data_dir,labels_file_name_digit)
    else:
        image_bin_file_path = os.path.join(data_dir,images_file_name_letter)
        label_bin_file_path = os.path.join(data_dir,labels_file_name_letter)


    image_bin_file = open(image_bin_file_path, 'rb')
    label_bin_file = open(label_bin_file_path, "rb")

    buffer = image_bin_file.read(16)

    magic, number, rows, cols = struct.unpack('iiii',buffer)

    image_bin_file.seek(16)

    buffer2 = image_bin_file.read(number*rows*cols)

    dat = struct.unpack(str(number*rows*cols)+'B',buffer2)

    image_dat = numpy.reshape(dat,[number,rows,cols])

    image_bin_file.close()

    buffer = label_bin_file.read(8)

    magic, number = struct.unpack('ii',buffer)

    label_bin_file.seek(8)

    buffer2 = label_bin_file.read(number)

    label_dat = struct.unpack(str(number) + 'B', buffer2)

    label_bin_file.close()

chars = list_subdir(image_dir)

digit , letter = read_all_images(chars)

write_file(digit,0)
write_file(letter,1)

# open_file()









